/*
 * Copyright (c) 2023, LXMS (lxinet.cn).
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package cn.lxinet.lxms.framework.core.utils;

import cn.lxinet.lxms.framework.core.exception.LxmsException;
import cn.lxinet.lxms.framework.core.message.CommonErrorCode;
import lombok.extern.slf4j.Slf4j;

import javax.crypto.Cipher;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

@Slf4j
public class RSAUtil {
    public static final String KEY_ALGORITHM = "RSA";
    private static final String PUBLIC_KEY = "RSAPublicKey";
    private static final String PRIVATE_KEY = "RSAPrivateKey";
    private Map<String, Object> keyMap = new HashMap<>(2);
    private static RSAUtil rsa;
    static {
        rsa = new RSAUtil();
        rsa.initKey(1024);
    }

    public static RSAUtil getInstance(){
        return rsa;
    }

    // 生成密钥对
    public void initKey(int keysize){
        KeyPairGenerator keyPairGen = null;
        try {
            keyPairGen = KeyPairGenerator.getInstance(KEY_ALGORITHM);
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        }
        // 设置密钥对的 bit 数，越大越安全
        keyPairGen.initialize(keysize);
        KeyPair keyPair = keyPairGen.generateKeyPair();

        // 获取公钥
        RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
        // 获取私钥
        RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
        keyMap.put(PUBLIC_KEY, publicKey);
        keyMap.put(PRIVATE_KEY, privateKey);
    }

    // 获取公钥字符串
    public String getPublicKeyStr() {
        // 获得 map 中的公钥对象，转为 key 对象
        Key key = (Key) keyMap.get(PUBLIC_KEY);
        // 编号返回字符串
        return encryptBASE64(key.getEncoded());
    }

    // 获取私钥字符串
    public String getPrivateKeyStr() {
        // 获得 map 中的私钥对象，转为 key 对象
        Key key = (Key) keyMap.get(PRIVATE_KEY);
        // 编号返回字符串
        return encryptBASE64(key.getEncoded());
    }

    // 获取公钥
    public PublicKey getPublicKey(String publicKeyString) throws NoSuchAlgorithmException, InvalidKeySpecException {
        byte[] publicKeyByte = Base64.getDecoder().decode(publicKeyString);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(publicKeyByte);
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        return keyFactory.generatePublic(keySpec);
    }

    // 获取私钥
    public PrivateKey getPrivateKey(String privateKeyString) throws Exception {
        byte[] privateKeyByte = Base64.getDecoder().decode(privateKeyString);
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKeyByte);
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        return keyFactory.generatePrivate(keySpec);
    }

    /**
     * BASE64 编号返回加密字符串
     *
     * @param key 需要编号的字节数组
     * @return 编号后的字符串
     */
    public String encryptBASE64(byte[] key) {
        return new String(Base64.getEncoder().encode(key));
    }

    /**
     * BASE64 解码，返回字节数组
     *
     * @param key 待解码的字符串
     * @return 解码后的字节数组
     */
    public byte[] decryptBASE64(String key) {
        return Base64.getDecoder().decode(key);
    }

    /**
     * 公钥加密
     *
     * @param text         待加密的明文字符串
     * @return 加密后的密文
     */
    public String encrypt(String text) {
        try {
            Cipher cipher = Cipher.getInstance(KEY_ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, getPublicKey(getPublicKeyStr()));
            byte[] tempBytes = cipher.doFinal(text.getBytes("UTF-8"));
            return Base64.getEncoder().encodeToString(tempBytes);
        } catch (Exception e) {
            log.error("加密字符串[" + text + "]时遇到异常", e);
            throw new LxmsException(CommonErrorCode.RSA_ENCRYPT_ERROR);
        }
    }

    /**
     * 私钥解密
     *
     * @param secretText    待解密的密文字符串
     * @return 解密后的明文
     */
    public String decrypt(String secretText) {
        try {
            // 生成私钥
            Cipher cipher = Cipher.getInstance(KEY_ALGORITHM);
            cipher.init(Cipher.DECRYPT_MODE, getPrivateKey(getPrivateKeyStr()));
            // 密文解码
            byte[] secretTextDecoded = Base64.getDecoder().decode(secretText.getBytes("UTF-8"));
            byte[] tempBytes = cipher.doFinal(secretTextDecoded);
            return new String(tempBytes);
        } catch (Exception e) {
            log.error("解密字符串[" + secretText + "]时遇到异常", e);
            throw new LxmsException(CommonErrorCode.RSA_DECRYPT_ERROR);
        }
    }

    public static void main(String[] args) {
        RSAUtil rsaUtil = RSAUtil.getInstance();
        String cipherText;
        // 原始明文
        String content = "lxms";

        String publicKey = rsaUtil.getPublicKeyStr();
        System.out.println("公钥:[" + publicKey + "]，长度:[" + publicKey.length() + "]");
        String privateKey = rsaUtil.getPrivateKeyStr();
        System.out.println("私钥:[" + privateKey + "]，长度:[" + privateKey.length() + "]");

        // 加密
        cipherText = rsaUtil.encrypt(content);
        System.out.println("加密后的密文:[" + cipherText + "]，长度:[" + cipherText.length() + "]");

        // 解密
        String plainText = rsaUtil.decrypt(cipherText);
        System.out.println("解密后明文:[" + plainText + "]");
    }
}
